import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
# 显示中文，Windows系统
plt.rcParams['font.sans-serif'] = 'SimHei'
# 支持符号
plt.rcParams['axes.unicode_minus'] = False

# 加载预训练的MNIST模型
def load_model():
    # 如果已有训练好的模型，可以直接加载
    # model = tf.keras.models.load_model('mnist_cnn_model.h5')
    # return model

    # 否则创建并训练一个新模型
    # 加载MNIST数据集
    (train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

    # 数据预处理
    train_images = train_images.reshape((60000, 28, 28, 1)).astype('float32') / 255
    test_images = test_images.reshape((10000, 28, 28, 1)).astype('float32') / 255

    train_labels = tf.keras.utils.to_categorical(train_labels)
    test_labels = tf.keras.utils.to_categorical(test_labels)

    # 构建CNN模型
    model = models.Sequential()
    model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))
    model.add(layers.MaxPooling2D((2, 2)))
    model.add(layers.Conv2D(64, (3, 3), activation='relu'))
    model.add(layers.MaxPooling2D((2, 2)))
    model.add(layers.Conv2D(64, (3, 3), activation='relu'))
    model.add(layers.Flatten())
    model.add(layers.Dense(64, activation='relu'))
    model.add(layers.Dense(10, activation='softmax'))

    # 编译模型
    model.compile(optimizer='adam',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    # 训练模型
    model.fit(train_images, train_labels, epochs=5, batch_size=64, validation_data=(test_images, test_labels))

    # 保存模型
    model.save('mnist_cnn_model.h5')

    return model


# 加载并预处理图片
def load_and_preprocess_image(image_path):
    # 打开图片
    img = Image.open(image_path)
    # 转换为灰度图
    img = img.convert('L')
    # 调整大小为28x28
    img = img.resize((28, 28))
    # 转换为numpy数组
    img_array = np.array(img)
    # 反转颜色（MNIST数据集是黑底白字）
    img_array = 255 - img_array
    # 归一化
    img_array = img_array.astype('float32') / 255
    # 重塑为模型输入格式 (1, 28, 28, 1)
    img_array = img_array.reshape(1, 28, 28, 1)
    return img, img_array


# 预测数字
def predict_digit(model, img_array):
    prediction = model.predict(img_array)
    digit = np.argmax(prediction)
    confidence = np.max(prediction) * 100
    return digit, confidence


# 主函数
def main():
    # 图片路径
    image_path =r'1.png'

    # 加载模型
    print("加载模型...")
    model = load_model()

    # 加载并预处理图片
    print("加载并预处理图片...")
    img, img_array = load_and_preprocess_image(image_path)

    # 预测数字
    print("进行预测...")
    digit, confidence = predict_digit(model, img_array)

    # 显示结果
    plt.figure(figsize=(10, 4))

    plt.subplot(1, 2, 1)
    plt.imshow(img, cmap='gray')
    plt.title('输入图片')
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.bar(range(10), model.predict(img_array)[0])
    plt.title('预测概率')
    plt.xticks(range(10))
    plt.xlabel('数字')
    plt.ylabel('概率')

    plt.tight_layout()
    plt.show()

    print(f"预测结果: {digit}")
    print(f"置信度: {confidence:.2f}%")


if __name__ == "__main__":
    main()
